# -*- coding: utf-8 -*-
'''
Created on 14.11.2019

@author: yu03
'''
import numpy as np
import os
import re
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit
import glob
from mpl_toolkits.mplot3d import Axes3D
from PyUeye_Unified.Hor_16line import folder_path, np_result_names, hor_index, hor_lines
import sys
from FFT_Interpolation import FFT_interpolation_nonlinearity_compare
from matplotlib import cm
from matplotlib.ticker import LinearLocator, FormatStrFormatter
from scipy import signal
from FFT_Interpolation import FFT_cal


# plt_show_mode = 'show' ### plt.show()
plt_show_mode = 'save' ### plt.save()


def Export_Data(out_name, header, out_str):
    print('Writing Data')
    with open(out_name,'w') as fid: ######################################################################################
        fid.writelines(header)
        fid.writelines(out_str)
    print('Finish Writing')

i = 0
fs_cam = 1000
''' 
    for all groups
''' 
for np_result in np_result_names:
    ''' 
        reading .npy file
    ''' 
    file_name = np_result.split('\\')[-1] ### x_lines.py
#     file_name = re.findall(r"openmode\\(.+).npy", np_result)[0]
    print(file_name)
    txt_name = folder_path + '\\' + file_name.split('.')[0] + '_results.txt'
    PSD_name = folder_path + '\\' + file_name.split('.')[0] + '_PSD.txt'
    i += 1
    f = open(np_result, 'rb')
    f_size = os.fstat(f.fileno()).st_size
    ''' 
        put 2-D results in "lines"
    ''' 
    lines = []
    time_sequence = []
    while f.tell() < f_size:
        line = np.load(f, allow_pickle=True)
        lines.append(line)
#         print(np.shape(line))
    print('%i: %s, %eB'%(i, file_name, f_size))
#     lines = np.array(lines)[::,::,300:2200]
#     lines = lines[2:4]
    lines = np.array(lines)
    print(np.shape(lines), 'Channle, lines, frames:')
    frame_num = np.shape(lines)[2]
    line_num = np.shape(lines)[1]
    
    hor_angle_set, hor_length_set = lines
    ''' 
        making average
    ''' 
    averaged_results = []
     
    hor_angle_avr = np.mean(hor_angle_set, axis=0)
    hor_length_avr = np.mean(hor_length_set, axis=0)
     
    averaged_results.append(hor_angle_avr)
    averaged_results.append(hor_length_avr)
    
    hor_angle_avr = hor_angle_avr - np.mean(hor_angle_avr) 
    hor_length_avr = hor_length_avr - np.mean(hor_length_avr)
    
    
    f_line, PS_hor_angle_avr = signal.welch(hor_angle_avr, fs_cam, nperseg=len(hor_angle_avr))
    PSD_hor_angle_avr = np.sqrt(PS_hor_angle_avr)
    f_line, PS_hor_length_avr = signal.welch(hor_length_avr, fs_cam, nperseg=len(hor_length_avr))
    PSD_hor_length_avr = np.sqrt(PS_hor_length_avr)

#     f_line, xxx, FFT_hor_angle_avr, xxx = FFT_cal(hor_angle_avr, 1/1000)
#     f_line, xxx, FFT_hor_length_avr, xxx = FFT_cal(hor_length_avr, 1/1000)
#     f_line, FFT_hor_angle_avr, FFT_hor_length_avr = f_line[1:], FFT_hor_angle_avr[1:], FFT_hor_length_avr[1:]

    ''' 
        Results output
    '''
    if 1:
        ''' 
            Header
        '''
        header = ['Fs = %e (Hz)\n' %fs_cam,#########################################################################################################
              'Channel_1: D\n',############################################################################################
              'Channel_2: Angle hor.\n',############################################################################################
              '-------------------------------------------------\n',
              ]
        ''' 
            Exporting TXT
        '''
        out_str = ['%f, %f\n' %(hor_length_avr[i], hor_angle_avr[i]) for i in range(len(hor_angle_avr))]
        Export_Data(txt_name, header, out_str)
      
        ''' 
            PSD Header
        '''
        header_FFT = ['Fs = %e (Hz)\n' %50,#########################################################################################################
              'Channel_1: PSD_f_line\n',############################################################################################
              'Channel_2: PSD_D\n',############################################################################################
              'Channel_3: PSD_hor.\n',############################################################################################
              '-------------------------------------------------\n',
              ]
        ''' 
            PSD Exporting TXT
        '''
        out_str_FFT = ['%f, %f, %f\n' %(f_line[i], PSD_hor_length_avr[i], PSD_hor_angle_avr[i]) for i in range(len(PSD_hor_length_avr))]
        Export_Data(PSD_name, header_FFT, out_str_FFT)
    
    ''' 
        Plotting averaged results
    '''   
    fig1 = plt.figure('Noise_PSD & PSD')
    plt.gcf().set_size_inches(18,9)
    
    ax1 = fig1.add_subplot(2, 2, 1)
    ax2 = fig1.add_subplot(2, 2, 2)
    ax3 = fig1.add_subplot(2, 2, 3)
    ax4 = fig1.add_subplot(2, 2, 4)
    
    ax1.plot(hor_angle_avr)
    ax2.plot(hor_length_avr)
    ax3.loglog(f_line, PSD_hor_angle_avr)
    ax4.loglog(f_line, PSD_hor_length_avr)
#     ax3.loglog(f_line, FFT_hor_angle_avr)
#     ax4.loglog(f_line, FFT_hor_length_avr)
    
    ax1.title.set_text('Hor. Tilt')
    ax2.title.set_text('Hor. Length')
    ax3.title.set_text('Hor. Tilt PSD')
    ax4.title.set_text('Hor. Length PSD')
    
    ax1.set_xlabel('Frame num.')
    ax2.set_xlabel('Frame num.')
    ax3.set_xlabel('Freq. \ Hz')
    ax4.set_xlabel('Freq. \ Hz')
    
    ax1.set_ylabel('Angle / urad')
    ax2.set_ylabel('Displacement / nm')
    ax3.set_ylabel('PSD / urad/sqrt(Hz)')
    ax4.set_ylabel('PSD / nm/sqrt(Hz)')
    
    ax1.grid(which='both', axis='both')
    ax2.grid(which='both', axis='both')
    ax3.grid(which='both', axis='both')
    ax4.grid(which='both', axis='both')
    
    ax3.set_ylim(1e-5, 1e-1)
    ax4.set_ylim(1e-5, 1)
    
    plt.tight_layout()
    if plt_show_mode == 'save':
        plt.savefig(folder_path+'\\'+file_name.split('.')[0]+'.png', dpi=300)
        plt.close()
    elif plt_show_mode == 'show':
        plt.show()
    else:
        sys.exit('Figure Save/Show Error:\n plt_show_mode = %s'%plt_show_mode)

    
#     ''' 
#         Plotting line results
#     '''   
#     y = [np.arange(frame_num)]*line_num
#     y = np.array(y)
#         
#     fig = plt.figure('hor_tilt')
#     x = [hor_lines]*frame_num
#     x = np.array(x).T
#     ax = fig.gca(projection='3d')
#     plt.gca().patch.set_facecolor('white')
#     for i in range(line_num):
#         ax.plot(x[i,:],y[i,:],hor_angle_set[i,:])
#      
#            
#     fig = plt.figure('hor_length')
#     x = [hor_lines]*frame_num
#     x = np.array(x).T
#     ax = fig.gca(projection='3d')
#     plt.gca().patch.set_facecolor('white')
#     for i in range(line_num):
#         ax.plot(x[i,:],y[i,:],hor_length_set[i,:])
#        
#     plt.show()


